{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/NeuromatchAcademy/course-content/blob/master/tutorials/W3D1_BayesianDecisions/student/W3D1_Tutorial1.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Neuromatch Academy: Week 3, Day 1, Tutorial 1\n",
    "# Bayes with a binary hidden state\n",
    "\n",
    "__Content creators:__ [insert your name here]\n",
    "\n",
    "__Content reviewers:__ "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial Objectives\n",
    "This is the first in a series of two core tutorials on Bayesian statistics. In these tutorials, we will explore the fundemental concepts of the Bayesian approach from two perspectives. This tutorial will work through an example of Bayesian inference and decision making using a binary hidden state. The second main tutorial extends these concepts to a continuous hidden state. In the next days, each of these basic ideas will be extended--first through time as we consider what happens when we infere a hidden state using multiple observations and when the hidden state changes across time. In the third day, we will introduce the notion of how to use inference and decisions to select actions for optimal control. For this tutorial, you will be introduced to our binary state fishing problem!\n",
    "\n",
    "This notebook will introduce the fundamental building blocks for Bayesian statistics: \n",
    "\n",
    "1. How do we use probability distributions to represent hidden states?\n",
    "2. How does marginalization work and how can we use it?\n",
    "3. How do we combine new information with our prior knowledge?\n",
    "4. How do we combine the possible loss (or gain) for making a decision with our probabilitic knowledge?\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 517
    },
    "execution": {
     "iopub.execute_input": "2021-06-02T02:11:13.107233Z",
     "iopub.status.busy": "2021-06-02T02:11:13.106602Z",
     "iopub.status.idle": "2021-06-02T02:11:13.202995Z",
     "shell.execute_reply": "2021-06-02T02:11:13.203403Z"
    },
    "outputId": "f2f0a0ff-8ebf-486d-8a45-c3e012298e7d"
   },
   "outputs": [],
   "source": [
    "#@title Video 1: Introduction to Bayesian Statistics\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id='JiEIn9QsrFg', width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup  \n",
    "Please execute the cells below to initialize the notebook environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-02T02:11:13.207441Z",
     "iopub.status.busy": "2021-06-02T02:11:13.206944Z",
     "iopub.status.idle": "2021-06-02T02:11:13.940550Z",
     "shell.execute_reply": "2021-06-02T02:11:13.940996Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import patches\n",
    "from matplotlib import transforms\n",
    "from matplotlib import gridspec\n",
    "from scipy.optimize import fsolve\n",
    "\n",
    "from collections import namedtuple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-02T02:11:13.945911Z",
     "iopub.status.busy": "2021-06-02T02:11:13.945396Z",
     "iopub.status.idle": "2021-06-02T02:11:14.111212Z",
     "shell.execute_reply": "2021-06-02T02:11:14.110711Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Figure Settings\n",
    "import ipywidgets as widgets       # interactive display\n",
    "from ipywidgets import GridspecLayout\n",
    "from IPython.display import clear_output\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "plt.style.use(\"https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/nma.mplstyle\")\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-02T02:11:14.129999Z",
     "iopub.status.busy": "2021-06-02T02:11:14.119358Z",
     "iopub.status.idle": "2021-06-02T02:11:14.175319Z",
     "shell.execute_reply": "2021-06-02T02:11:14.174621Z"
    }
   },
   "outputs": [],
   "source": [
    "# @title Plotting Functions\n",
    "\n",
    "def plot_joint_probs(P, ):\n",
    "    assert np.all(P >= 0), \"probabilities should be >= 0\"\n",
    "    # normalize if not\n",
    "    P = P / np.sum(P)\n",
    "    marginal_y = np.sum(P,axis=1)\n",
    "    marginal_x = np.sum(P,axis=0)\n",
    "\n",
    "    # definitions for the axes\n",
    "    left, width = 0.1, 0.65\n",
    "    bottom, height = 0.1, 0.65\n",
    "    spacing = 0.005\n",
    "\n",
    "    # start with a square Figure\n",
    "    fig = plt.figure(figsize=(5, 5))\n",
    "\n",
    "    joint_prob = [left, bottom, width, height]\n",
    "    rect_histx = [left, bottom + height + spacing, width, 0.2]\n",
    "    rect_histy = [left + width + spacing, bottom, 0.2, height]\n",
    "\n",
    "    rect_x_cmap = plt.cm.Blues\n",
    "    rect_y_cmap = plt.cm.Reds\n",
    "\n",
    "    # Show joint probs and marginals\n",
    "    ax = fig.add_axes(joint_prob)\n",
    "    ax_x = fig.add_axes(rect_histx, sharex=ax)\n",
    "    ax_y = fig.add_axes(rect_histy, sharey=ax)\n",
    "\n",
    "    # Show joint probs and marginals\n",
    "    ax.matshow(P,vmin=0., vmax=1., cmap='Greys')\n",
    "    ax_x.bar(0, marginal_x[0], facecolor=rect_x_cmap(marginal_x[0]))\n",
    "    ax_x.bar(1, marginal_x[1], facecolor=rect_x_cmap(marginal_x[1]))\n",
    "    ax_y.barh(0, marginal_y[0], facecolor=rect_y_cmap(marginal_y[0]))\n",
    "    ax_y.barh(1, marginal_y[1], facecolor=rect_y_cmap(marginal_y[1]))\n",
    "    # set limits\n",
    "    ax_x.set_ylim([0,1])\n",
    "    ax_y.set_xlim([0,1])\n",
    "\n",
    "    # show values\n",
    "    ind = np.arange(2)\n",
    "    x,y = np.meshgrid(ind,ind)\n",
    "    for i,j in zip(x.flatten(), y.flatten()):\n",
    "        c = f\"{P[i,j]:.2f}\"\n",
    "        ax.text(j,i, c, va='center', ha='center', color='black')\n",
    "    for i in ind:\n",
    "        v = marginal_x[i]\n",
    "        c = f\"{v:.2f}\"\n",
    "        ax_x.text(i, v +0.1, c, va='center', ha='center', color='black')\n",
    "        v = marginal_y[i]\n",
    "        c = f\"{v:.2f}\"\n",
    "        ax_y.text(v+0.2, i, c, va='center', ha='center', color='black')\n",
    "\n",
    "    # set up labels\n",
    "    ax.xaxis.tick_bottom()\n",
    "    ax.yaxis.tick_left()\n",
    "    ax.set_xticks([0,1])\n",
    "    ax.set_yticks([0,1])\n",
    "    ax.set_xticklabels(['Silver','Gold'])\n",
    "    ax.set_yticklabels(['Small', 'Large'])\n",
    "    ax.set_xlabel('color')\n",
    "    ax.set_ylabel('size')\n",
    "    ax_x.axis('off')\n",
    "    ax_y.axis('off')\n",
    "    return fig\n",
    "# test\n",
    "# P = np.random.rand(2,2)\n",
    "# P = np.asarray([[0.9, 0.8], [0.4, 0.1]])\n",
    "# P = P / np.sum(P)\n",
    "# fig = plot_joint_probs(P)\n",
    "# plt.show(fig)\n",
    "# plt.close(fig)\n",
    "\n",
    "\n",
    "# fig = plot_prior_likelihood(0.5, 0.3)\n",
    "# plt.show(fig)\n",
    "# plt.close(fig)\n",
    "\n",
    "\n",
    "def plot_prior_likelihood_posterior(prior, likelihood, posterior):\n",
    "\n",
    "    # definitions for the axes\n",
    "    left, width = 0.05, 0.3\n",
    "    bottom, height = 0.05, 0.9\n",
    "    padding = 0.1\n",
    "    small_width = 0.1\n",
    "    left_space = left + small_width + padding\n",
    "    added_space = padding + width\n",
    "\n",
    "    fig = plt.figure(figsize=(10, 4))\n",
    "\n",
    "    rect_prior = [left, bottom, small_width, height]\n",
    "    rect_likelihood = [left_space , bottom , width, height]\n",
    "    rect_posterior = [left_space +  added_space, bottom , width, height]\n",
    "\n",
    "    ax_prior = fig.add_axes(rect_prior)\n",
    "    ax_likelihood = fig.add_axes(rect_likelihood, sharey=ax_prior)\n",
    "    ax_posterior = fig.add_axes(rect_posterior, sharey = ax_prior)\n",
    "\n",
    "    rect_colormap = plt.cm.Blues\n",
    "\n",
    "    # Show posterior probs and marginals\n",
    "    ax_prior.barh(0, prior[0], facecolor = rect_colormap(prior[0, 0]))\n",
    "    ax_prior.barh(1, prior[1], facecolor = rect_colormap(prior[1, 0]))\n",
    "    ax_likelihood.matshow(likelihood, vmin=0., vmax=1., cmap='Reds')\n",
    "    ax_posterior.matshow(posterior, vmin=0., vmax=1., cmap='Greens')\n",
    "\n",
    "\n",
    "    # Probabilities plot details\n",
    "    ax_prior.set(xlim = [1, 0], yticks = [0, 1], yticklabels = ['left', 'right'],\n",
    "                 ylabel = 'state (s)', title = \"Prior p(s)\")\n",
    "    ax_prior.axis('off')\n",
    "\n",
    "    # Likelihood plot details\n",
    "    ax_likelihood.set(xticks = [0, 1], xticklabels = ['fish', 'no fish'],\n",
    "                  yticks = [0, 1], yticklabels = ['left', 'right'],\n",
    "                   ylabel = 'state (s)', xlabel = 'measurement (m)',\n",
    "                   title = 'Likelihood p(m (right) | s)')\n",
    "    ax_likelihood.xaxis.set_ticks_position('bottom')\n",
    "    ax_likelihood.spines['left'].set_visible(False)\n",
    "    ax_likelihood.spines['bottom'].set_visible(False)\n",
    "\n",
    "    # Posterior plot details\n",
    "\n",
    "    ax_posterior.set(xticks = [0, 1], xticklabels = ['fish', 'no fish'],\n",
    "                  yticks = [0, 1], yticklabels = ['left', 'right'],\n",
    "                   ylabel = 'state (s)', xlabel = 'measurement (m)',\n",
    "                   title = 'Posterior p(s | m)')\n",
    "    ax_posterior.xaxis.set_ticks_position('bottom')\n",
    "    ax_posterior.spines['left'].set_visible(False)\n",
    "    ax_posterior.spines['bottom'].set_visible(False)\n",
    "\n",
    "\n",
    "    # show values\n",
    "    ind = np.arange(2)\n",
    "    x,y = np.meshgrid(ind,ind)\n",
    "    for i,j in zip(x.flatten(), y.flatten()):\n",
    "        c = f\"{posterior[i,j]:.2f}\"\n",
    "        ax_posterior.text(j,i, c, va='center', ha='center', color='black')\n",
    "    for i,j in zip(x.flatten(), y.flatten()):\n",
    "        c = f\"{likelihood[i,j]:.2f}\"\n",
    "        ax_likelihood.text(j,i, c, va='center', ha='center', color='black')\n",
    "    for i in ind:\n",
    "        v = prior[i, 0]\n",
    "        c = f\"{v:.2f}\"\n",
    "        ax_prior.text(v+0.2, i, c, va='center', ha='center', color='black')\n",
    "\n",
    "\n",
    "def plot_prior_likelihood(ps, p_a_s1, p_a_s0, measurement):\n",
    "    likelihood = np.asarray([[p_a_s1, 1-p_a_s1],[p_a_s0, 1-p_a_s0]])\n",
    "    assert 0.0 <= ps <= 1.0\n",
    "    prior = np.asarray([ps, 1 - ps])\n",
    "    if measurement:\n",
    "        posterior = likelihood[:, 0] * prior\n",
    "    else:\n",
    "        posterior = (likelihood[:, 1] * prior).reshape(-1)\n",
    "    posterior /= np.sum(posterior)\n",
    "\n",
    "    # definitions for the axes\n",
    "    left, width = 0.05, 0.3\n",
    "    bottom, height = 0.05, 0.9\n",
    "    padding = 0.1\n",
    "    small_width = 0.22\n",
    "    left_space = left + small_width + padding\n",
    "    small_padding = 0.05\n",
    "\n",
    "    fig = plt.figure(figsize=(10, 4))\n",
    "\n",
    "    rect_prior = [left, bottom, small_width, height]\n",
    "    rect_likelihood = [left_space , bottom , width, height]\n",
    "    rect_posterior = [left_space + width + small_padding, bottom , small_width, height]\n",
    "\n",
    "    ax_prior = fig.add_axes(rect_prior)\n",
    "    ax_likelihood = fig.add_axes(rect_likelihood, sharey=ax_prior)\n",
    "    ax_posterior = fig.add_axes(rect_posterior, sharey=ax_prior)\n",
    "\n",
    "    prior_colormap = plt.cm.Blues\n",
    "    posterior_colormap = plt.cm.Greens\n",
    "\n",
    "    # Show posterior probs and marginals\n",
    "    ax_prior.barh(0, prior[0], facecolor = prior_colormap(prior[0]))\n",
    "    ax_prior.barh(1, prior[1], facecolor = prior_colormap(prior[1]))\n",
    "    ax_likelihood.matshow(likelihood, vmin=0., vmax=1., cmap='Reds')\n",
    "    # ax_posterior.matshow(posterior, vmin=0., vmax=1., cmap='')\n",
    "    ax_posterior.barh(0, posterior[0], facecolor = posterior_colormap(posterior[0]))\n",
    "    ax_posterior.barh(1, posterior[1], facecolor = posterior_colormap(posterior[1]))\n",
    "\n",
    "    # Probabilities plot details\n",
    "    ax_prior.set(xlim = [1, 0], yticks = [0, 1], yticklabels = ['left', 'right'],\n",
    "                 ylabel = 'state (s)', title = \"Prior p(s)\")\n",
    "    ax_prior.axis('off')\n",
    "\n",
    "    # Likelihood plot details\n",
    "    ax_likelihood.set(xticks = [0, 1], xticklabels = ['fish', 'no fish'],\n",
    "                  yticks = [0, 1], yticklabels = ['left', 'right'],\n",
    "                   ylabel = 'state (s)', xlabel = 'measurement (m)',\n",
    "                   title = 'Likelihood p(m | s)')\n",
    "    ax_likelihood.xaxis.set_ticks_position('bottom')\n",
    "    ax_likelihood.spines['left'].set_visible(False)\n",
    "    ax_likelihood.spines['bottom'].set_visible(False)\n",
    "\n",
    "    # Posterior plot details\n",
    "    ax_posterior.set(xlim = [0, 1], yticks = [0, 1], yticklabels = ['left', 'right'],\n",
    "                 ylabel = 'state (s)', title = \"Posterior p(s | m)\")\n",
    "    ax_posterior.axis('off')\n",
    "    # ax_posterior.set(xticks = [0, 1], xticklabels = ['fish', 'no fish'],\n",
    "    #               yticks = [0, 1], yticklabels = ['left', 'right'],\n",
    "    #                ylabel = 'state (s)', xlabel = 'measurement (m)',\n",
    "    #                title = 'Posterior p(s | m)')\n",
    "    # ax_posterior.xaxis.set_ticks_position('bottom')\n",
    "    # ax_posterior.spines['left'].set_visible(False)\n",
    "    # ax_posterior.spines['bottom'].set_visible(False)\n",
    "\n",
    "\n",
    "    # show values\n",
    "    ind = np.arange(2)\n",
    "    x,y = np.meshgrid(ind,ind)\n",
    "    # for i,j in zip(x.flatten(), y.flatten()):\n",
    "    #     c = f\"{posterior[i,j]:.2f}\"\n",
    "    #     ax_posterior.text(j,i, c, va='center', ha='center', color='black')\n",
    "    for i in ind:\n",
    "        v = posterior[i]\n",
    "        c = f\"{v:.2f}\"\n",
    "        ax_posterior.text(v+0.2, i, c, va='center', ha='center', color='black')\n",
    "    for i,j in zip(x.flatten(), y.flatten()):\n",
    "        c = f\"{likelihood[i,j]:.2f}\"\n",
    "        ax_likelihood.text(j,i, c, va='center', ha='center', color='black')\n",
    "    for i in ind:\n",
    "        v = prior[i]\n",
    "        c = f\"{v:.2f}\"\n",
    "        ax_prior.text(v+0.2, i, c, va='center', ha='center', color='black')\n",
    "    return fig\n",
    "\n",
    "# fig = plot_prior_likelihood(0.5, 0.3)\n",
    "# plt.show(fig)\n",
    "# plt.close(fig)\n",
    "\n",
    "\n",
    "from matplotlib import colors\n",
    "def plot_utility(ps):\n",
    "    prior = np.asarray([ps, 1 - ps])\n",
    "\n",
    "    utility = np.array([[2, -3], [-2, 1]])\n",
    "\n",
    "    expected = prior @ utility\n",
    "\n",
    "    # definitions for the axes\n",
    "    left, width = 0.05, 0.16\n",
    "    bottom, height = 0.05, 0.9\n",
    "    padding = 0.04\n",
    "    small_width = 0.1\n",
    "    left_space = left + small_width + padding\n",
    "    added_space = padding + width\n",
    "\n",
    "    fig = plt.figure(figsize=(17, 3))\n",
    "\n",
    "    rect_prior = [left, bottom, small_width, height]\n",
    "    rect_utility = [left + added_space , bottom , width, height]\n",
    "    rect_expected = [left + 2* added_space, bottom , width, height]\n",
    "\n",
    "    ax_prior = fig.add_axes(rect_prior)\n",
    "    ax_utility = fig.add_axes(rect_utility, sharey=ax_prior)\n",
    "    ax_expected = fig.add_axes(rect_expected)\n",
    "\n",
    "    rect_colormap = plt.cm.Blues\n",
    "\n",
    "    # Data of plots\n",
    "    ax_prior.barh(0, prior[0], facecolor = rect_colormap(prior[0]))\n",
    "    ax_prior.barh(1, prior[1], facecolor = rect_colormap(prior[1]))\n",
    "    ax_utility.matshow(utility, cmap='cool')\n",
    "    norm = colors.Normalize(vmin=-3, vmax=3)\n",
    "    ax_expected.bar(0, expected[0], facecolor = rect_colormap(norm(expected[0])))\n",
    "    ax_expected.bar(1, expected[1], facecolor = rect_colormap(norm(expected[1])))\n",
    "\n",
    "    # Probabilities plot details\n",
    "    ax_prior.set(xlim = [1, 0], yticks = [0, 1], yticklabels = ['left', 'right'],\n",
    "                 ylabel = 'state (s)', title = \"Probability of state\")\n",
    "    ax_prior.axis('off')\n",
    "\n",
    "    # Utility plot details\n",
    "    ax_utility.set(xticks = [0, 1], xticklabels = ['left', 'right'],\n",
    "                  yticks = [0, 1], yticklabels = ['left', 'right'],\n",
    "                   ylabel = 'state (s)', xlabel = 'action (a)',\n",
    "                   title = 'Utility')\n",
    "    ax_utility.xaxis.set_ticks_position('bottom')\n",
    "    ax_utility.spines['left'].set_visible(False)\n",
    "    ax_utility.spines['bottom'].set_visible(False)\n",
    "\n",
    "    # Expected utility plot details\n",
    "    ax_expected.set(title = 'Expected utility', ylim = [-3, 3],\n",
    "                    xticks = [0, 1], xticklabels = ['left', 'right'],\n",
    "                    xlabel = 'action (a)',\n",
    "                    yticks = [])\n",
    "    ax_expected.xaxis.set_ticks_position('bottom')\n",
    "    ax_expected.spines['left'].set_visible(False)\n",
    "    ax_expected.spines['bottom'].set_visible(False)\n",
    "\n",
    "    # show values\n",
    "    ind = np.arange(2)\n",
    "    x,y = np.meshgrid(ind,ind)\n",
    "\n",
    "    for i,j in zip(x.flatten(), y.flatten()):\n",
    "        c = f\"{utility[i,j]:.2f}\"\n",
    "        ax_utility.text(j,i, c, va='center', ha='center', color='black')\n",
    "    for i in ind:\n",
    "        v = prior[i]\n",
    "        c = f\"{v:.2f}\"\n",
    "        ax_prior.text(v+0.2, i, c, va='center', ha='center', color='black')\n",
    "    for i in ind:\n",
    "        v = expected[i]\n",
    "        c = f\"{v:.2f}\"\n",
    "        ax_expected.text(i, 2.5, c, va='center', ha='center', color='black')\n",
    "\n",
    "    return fig\n",
    "\n",
    "\n",
    "def plot_prior_likelihood_utility(ps, p_a_s1, p_a_s0,measurement):\n",
    "    assert 0.0 <= ps <= 1.0\n",
    "    assert 0.0 <= p_a_s1 <= 1.0\n",
    "    assert 0.0 <= p_a_s0 <= 1.0\n",
    "    prior = np.asarray([ps, 1 - ps])\n",
    "    likelihood = np.asarray([[p_a_s1, 1-p_a_s1],[p_a_s0, 1-p_a_s0]])\n",
    "    utility = np.array([[2.0, -3.0], [-2.0, 1.0]])\n",
    "    # expected = np.zeros_like(utility)\n",
    "\n",
    "    if measurement:\n",
    "        posterior = likelihood[:, 0] * prior\n",
    "    else:\n",
    "        posterior = (likelihood[:, 1] * prior).reshape(-1)\n",
    "    posterior /= np.sum(posterior)\n",
    "    # expected[:, 0] = utility[:, 0] * posterior\n",
    "    # expected[:, 1] = utility[:, 1] * posterior\n",
    "    expected = posterior @ utility\n",
    "\n",
    "    # definitions for the axes\n",
    "    left, width = 0.05, 0.15\n",
    "    bottom, height = 0.05, 0.9\n",
    "    padding = 0.05\n",
    "    small_width = 0.1\n",
    "    large_padding = 0.07\n",
    "    left_space = left + small_width + large_padding\n",
    "\n",
    "    fig = plt.figure(figsize=(17, 4))\n",
    "\n",
    "    rect_prior = [left, bottom+0.05, small_width, height-0.1]\n",
    "    rect_likelihood = [left_space, bottom , width, height]\n",
    "    rect_posterior = [left_space + padding + width - 0.02, bottom+0.05 , small_width, height-0.1]\n",
    "    rect_utility = [left_space + padding + width + padding + small_width, bottom , width, height]\n",
    "    rect_expected = [left_space + padding + width + padding + small_width + padding + width, bottom+0.05 , width, height-0.1]\n",
    "\n",
    "    ax_likelihood = fig.add_axes(rect_likelihood)\n",
    "    ax_prior = fig.add_axes(rect_prior, sharey=ax_likelihood)\n",
    "    ax_posterior = fig.add_axes(rect_posterior, sharey=ax_likelihood)\n",
    "    ax_utility = fig.add_axes(rect_utility, sharey=ax_posterior)\n",
    "    ax_expected = fig.add_axes(rect_expected)\n",
    "\n",
    "\n",
    "    prior_colormap = plt.cm.Blues\n",
    "    posterior_colormap = plt.cm.Greens\n",
    "    expected_colormap = plt.cm.Wistia\n",
    "\n",
    "    # Show posterior probs and marginals\n",
    "    ax_prior.barh(0, prior[0], facecolor = prior_colormap(prior[0]))\n",
    "    ax_prior.barh(1, prior[1], facecolor = prior_colormap(prior[1]))\n",
    "    ax_likelihood.matshow(likelihood, vmin=0., vmax=1., cmap='Reds')\n",
    "    ax_posterior.barh(0, posterior[0], facecolor = posterior_colormap(posterior[0]))\n",
    "    ax_posterior.barh(1, posterior[1], facecolor = posterior_colormap(posterior[1]))\n",
    "    ax_utility.matshow(utility, vmin=0., vmax=1., cmap='cool')\n",
    "    # ax_expected.matshow(expected, vmin=0., vmax=1., cmap='Wistia')\n",
    "    ax_expected.bar(0, expected[0], facecolor = expected_colormap(expected[0]))\n",
    "    ax_expected.bar(1, expected[1], facecolor = expected_colormap(expected[1]))\n",
    "\n",
    "    # Probabilities plot details\n",
    "    ax_prior.set(xlim = [1, 0], yticks = [0, 1], yticklabels = ['left', 'right'],\n",
    "                 ylabel = 'state (s)', title = \"Prior p(s)\")\n",
    "    ax_prior.axis('off')\n",
    "\n",
    "    # Likelihood plot details\n",
    "    ax_likelihood.set(xticks = [0, 1], xticklabels = ['fish', 'no fish'],\n",
    "                  yticks = [0, 1], yticklabels = ['left', 'right'],\n",
    "                   ylabel = 'state (s)', xlabel = 'measurement (m)',\n",
    "                   title = 'Likelihood p(m | s)')\n",
    "    ax_likelihood.xaxis.set_ticks_position('bottom')\n",
    "    ax_likelihood.spines['left'].set_visible(False)\n",
    "    ax_likelihood.spines['bottom'].set_visible(False)\n",
    "\n",
    "    # Posterior plot details\n",
    "    ax_posterior.set(xlim = [0, 1], yticks = [0, 1], yticklabels = ['left', 'right'],\n",
    "                 ylabel = 'state (s)', title = \"Posterior p(s | m)\")\n",
    "    ax_posterior.axis('off')\n",
    "\n",
    "    # Utility plot details\n",
    "    ax_utility.set(xticks = [0, 1], xticklabels = ['left', 'right'],\n",
    "                   xlabel = 'action (a)',\n",
    "                   title = 'Utility')\n",
    "    ax_utility.xaxis.set_ticks_position('bottom')\n",
    "    ax_utility.spines['left'].set_visible(False)\n",
    "    ax_utility.spines['bottom'].set_visible(False)\n",
    "\n",
    "    # Expected Utility plot details\n",
    "    ax_expected.set(ylim = [-2, 2], xticks = [0, 1], xticklabels = ['left', 'right'],\n",
    "                 xlabel = 'action (a)', title = 'Expected utility', yticks=[])\n",
    "    # ax_expected.axis('off')\n",
    "    ax_expected.spines['left'].set_visible(False)\n",
    "    # ax_expected.set(xticks = [0, 1], xticklabels = ['left', 'right'],\n",
    "    #                 xlabel = 'action (a)',\n",
    "    #                title = 'Expected utility')\n",
    "    # ax_expected.xaxis.set_ticks_position('bottom')\n",
    "    # ax_expected.spines['left'].set_visible(False)\n",
    "    # ax_expected.spines['bottom'].set_visible(False)\n",
    "\n",
    "    # show values\n",
    "    ind = np.arange(2)\n",
    "    x,y = np.meshgrid(ind,ind)\n",
    "    for i in ind:\n",
    "        v = posterior[i]\n",
    "        c = f\"{v:.2f}\"\n",
    "        ax_posterior.text(v+0.2, i, c, va='center', ha='center', color='black')\n",
    "    for i,j in zip(x.flatten(), y.flatten()):\n",
    "        c = f\"{likelihood[i,j]:.2f}\"\n",
    "        ax_likelihood.text(j,i, c, va='center', ha='center', color='black')\n",
    "    for i,j in zip(x.flatten(), y.flatten()):\n",
    "        c = f\"{utility[i,j]:.2f}\"\n",
    "        ax_utility.text(j,i, c, va='center', ha='center', color='black')\n",
    "    # for i,j in zip(x.flatten(), y.flatten()):\n",
    "    #     c = f\"{expected[i,j]:.2f}\"\n",
    "    #     ax_expected.text(j,i, c, va='center', ha='center', color='black')\n",
    "    for i in ind:\n",
    "        v = prior[i]\n",
    "        c = f\"{v:.2f}\"\n",
    "        ax_prior.text(v+0.2, i, c, va='center', ha='center', color='black')\n",
    "    for i in ind:\n",
    "        v = expected[i]\n",
    "        c = f\"{v:.2f}\"\n",
    "        ax_expected.text(i, v, c, va='center', ha='center', color='black')\n",
    "\n",
    "    # # show values\n",
    "    # ind = np.arange(2)\n",
    "    # x,y = np.meshgrid(ind,ind)\n",
    "    # for i,j in zip(x.flatten(), y.flatten()):\n",
    "    #     c = f\"{P[i,j]:.2f}\"\n",
    "    #     ax.text(j,i, c, va='center', ha='center', color='white')\n",
    "    # for i in ind:\n",
    "    #     v = marginal_x[i]\n",
    "    #     c = f\"{v:.2f}\"\n",
    "    #     ax_x.text(i, v +0.2, c, va='center', ha='center', color='black')\n",
    "    #     v = marginal_y[i]\n",
    "    #     c = f\"{v:.2f}\"\n",
    "    #     ax_y.text(v+0.2, i, c, va='center', ha='center', color='black')\n",
    "\n",
    "    return fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-02T02:11:14.182747Z",
     "iopub.status.busy": "2021-06-02T02:11:14.181503Z",
     "iopub.status.idle": "2021-06-02T02:11:14.183330Z",
     "shell.execute_reply": "2021-06-02T02:11:14.183745Z"
    }
   },
   "outputs": [],
   "source": [
    "# @title Helper Functions\n",
    "\n",
    "def compute_marginal(px, py, cor):\n",
    "    # calculate 2x2 joint probabilities given marginals p(x=1), p(y=1) and correlation\n",
    "    p11 = px*py + cor*np.sqrt(px*py*(1-px)*(1-py))\n",
    "    p01 = px - p11\n",
    "    p10 = py - p11\n",
    "    p00 = 1.0 - p11 - p01 - p10\n",
    "    return np.asarray([[p00, p01], [p10, p11]])\n",
    "# test\n",
    "# print(compute_marginal(0.4, 0.6, -0.8))\n",
    "\n",
    "\n",
    "def compute_cor_range(px,py):\n",
    "    # Calculate the allowed range of correlation values given marginals p(x=1) and p(y=1)\n",
    "    def p11(corr):\n",
    "        return px*py + corr*np.sqrt(px*py*(1-px)*(1-py))\n",
    "    def p01(corr):\n",
    "        return px - p11(corr)\n",
    "    def p10(corr):\n",
    "        return py - p11(corr)\n",
    "    def p00(corr):\n",
    "        return 1.0 - p11(corr) - p01(corr) - p10(corr)\n",
    "    Cmax = min(fsolve(p01, 0.0), fsolve(p10, 0.0))\n",
    "    Cmin = max(fsolve(p11, 0.0), fsolve(p00, 0.0))\n",
    "    return Cmin, Cmax\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Section 1: Gone Fishin'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 517
    },
    "execution": {
     "iopub.execute_input": "2021-06-02T02:11:14.188454Z",
     "iopub.status.busy": "2021-06-02T02:11:14.187979Z",
     "iopub.status.idle": "2021-06-02T02:11:14.254337Z",
     "shell.execute_reply": "2021-06-02T02:11:14.254793Z"
    },
    "outputId": "2209f6d0-af50-45f2-fc41-3117db3246dd"
   },
   "outputs": [],
   "source": [
    "#@title Video 2: Gone Fishin'\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id='McALsTzb494', width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You were just introduced to the **binary hidden state problem** we are going to explore. You need to decide which side to fish on. We know fish like to school together. On different days the school of fish is either on the left or right side, but we don’t know what the case is today. We will represent our knowledge probabilistically, asking how to make a decision (where to decide the fish are or where to fish) and what to expect in terms of gains or losses. In the next two sections we will consider just the probability of where the fish might be and what you gain or lose by choosing where to fish.\n",
    "\n",
    "Remember, you can either think of your self as a scientist conducting an experiment or as a brain trying to make a decision. The Bayesian approach is the same!\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Section 2: Deciding where to fish \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 517
    },
    "execution": {
     "iopub.execute_input": "2021-06-02T02:11:14.259686Z",
     "iopub.status.busy": "2021-06-02T02:11:14.259202Z",
     "iopub.status.idle": "2021-06-02T02:11:14.323724Z",
     "shell.execute_reply": "2021-06-02T02:11:14.323229Z"
    },
    "outputId": "5fe78511-352c-4122-80d3-6d8ad232cc95"
   },
   "outputs": [],
   "source": [
    "#@title Video 3: Utility\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id='xvIVZrqF_5s', width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You know the probability that the school of fish is on the left side of the dock today, $P(s = left)$. You also know the probability that it is on the right, $P(s = right)$, because these two probabilities must add up to 1. You need to decide where to fish. It may seem obvious - you could just fish on the side where the probability of the fish being is higher! Unfortunately, decisions and actions are always a little more complicated. Deciding to fish may be influenced by more than just the probability of the school of fish being there as we saw by the potential issues of submarines and sunburn. \n",
    "\n",
    "We quantify these factors numerically using **utility**, which describes the consequences of your actions: how much value you gain (or if negative, lose) given the state of the world ($s$) and the action you take ($a$). In our example, our utility can be summarized as:\n",
    "\n",
    "| Utility: U(s,a)   | a = left   | a = right  |\n",
    "| ----------------- |----------|----------|\n",
    "| s = Left          | 2          | -3         |\n",
    "| s = right         | -2         |  1         |\n",
    "\n",
    "To use utility to choose an action, we calculate the **expected utility** of that action by weighing these utilities with the probability of that state occuring. This allows us to choose actions by taking probabilities of events into account: we don't care if the outcome of an action-state pair is a loss if the probability of that state is very low. We can formalize this as:\n",
    "\n",
    "$$\\text{Expected utility of action a} = \\sum_{s}U(s,a)P(s) $$\n",
    "\n",
    "In other words, the expected utility of an action a is the sum over possible states of the utility of that action and state times the probability of that state.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interactive Demo 2: Exploring the decision\n",
    "\n",
    "Let's start to get a sense of how all this works. \n",
    "\n",
    "Take a look at the interactive demo below. You can change the probability that the school of fish is on the left side ($p(s = left)$ using the slider. You will see the utility matrix and the corresponding expected utility of each action.\n",
    "\n",
    "First, make sure you understand how the expected utility of each action is being computed from the probabilities and the utility values. In the initial state: the probability of the fish being on the left is 0.9 and on the right is 0.1. The expected utility of the action of fishing on the left is then $U(s = left,a = left)p(s = left) + U(s = right,a = left)p(s = right) = 2(0.9) + -2(0.1) = 1.6$.\n",
    "\n",
    "For each of these scenarios, think and discuss first. Then use the demo to try out each and see if your action would have been correct (that is, if the expected value of that action is the highest).\n",
    "\n",
    "\n",
    "1.  You just arrived at the dock for the first time and have no sense of where the fish might be. So you guess that the probability of the school being on the left side is 0.5 (so the probability on the right side is also 0.5). Which side would you choose to fish on given our utility values?\n",
    "2.  You think that the probability of the school being on the left side is very low (0.1) and correspondingly high on the right side (0.9). Which side would you choose to fish on given our utility values?\n",
    "3.  What would you choose if the probability of the school being on the left side is slightly lower than on the right side (0. 4 vs 0.6)?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-02T02:11:14.342580Z",
     "iopub.status.busy": "2021-06-02T02:11:14.335893Z",
     "iopub.status.idle": "2021-06-02T02:11:14.561208Z",
     "shell.execute_reply": "2021-06-02T02:11:14.561589Z"
    }
   },
   "outputs": [],
   "source": [
    "# @markdown Execute this cell to use the widget\n",
    "ps_widget = widgets.FloatSlider(0.9, description='p(s = left)', min=0.0, max=1.0, step=0.01)\n",
    "\n",
    "@widgets.interact(\n",
    "    ps = ps_widget,\n",
    ")\n",
    "def make_utility_plot(ps):\n",
    "    fig = plot_utility(ps)\n",
    "    plt.show(fig)\n",
    "    plt.close(fig)\n",
    "    return None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "execution": {
     "iopub.execute_input": "2021-06-02T02:11:14.566668Z",
     "iopub.status.busy": "2021-06-02T02:11:14.565543Z",
     "iopub.status.idle": "2021-06-02T02:11:14.567222Z",
     "shell.execute_reply": "2021-06-02T02:11:14.567614Z"
    }
   },
   "source": [
    "[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D1_BayesianDecisions/solutions/W3D1_Tutorial1_Solution_459cbf35.py)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, you have seen that both the utility of various state and action pairs and our knowledge of the probability of each state affects your decision. Importantly, we want our knowledge of the probability of each state to be as accurate as possible! \n",
    "\n",
    "So how do we know these probabilities? We may have prior knowledge from years of fishing at the same dock. Over those years, we may have learned that the fish are more likely to be on the left side for example. We want to make sure this knowledge is as accurate as possible though. To do this, we want to collect more data, or take some more measurements! For the next few sections, we will focus on making our knowledge of the probability as accurate as possible, before coming back to using utility to make decisions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Section 3: Likelihood of the fish being on either side\n",
    " \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 517
    },
    "execution": {
     "iopub.execute_input": "2021-06-02T02:11:14.572388Z",
     "iopub.status.busy": "2021-06-02T02:11:14.571892Z",
     "iopub.status.idle": "2021-06-02T02:11:14.643657Z",
     "shell.execute_reply": "2021-06-02T02:11:14.644055Z"
    },
    "outputId": "ce7995af-5602-4cb3-b403-3f759d1bb0af"
   },
   "outputs": [],
   "source": [
    "#@title Video 4: Likelihood\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id='l4m0JzMWGio', width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we'll think about what it means to take a measurement (also often called an observation or just data) and what it tells you about what the hidden state may be. Specifically, we'll be looking at the **likelihood**, which is the probability of your measurement ($m$) given the hidden state ($s$): $P(m | s)$. Remember that in this case, the hidden state is which side of the dock the school of fish is on.\n",
    "\n",
    "We will watch someone fish (for let's say 10 minutes) and our measurement is whether they catch a fish or not. We know something about what catching a fish means for the likelihood of the fish being on one side or the other. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Think! 3: Guessing the location of the fish\n",
    "\n",
    "Let's say we go to different dock from the one in the video. Here, there are different probabilities of catching fish given the state of the world. In this case, if they fish on the side of the dock where the fish are, they have a 70% chance of catching a fish. Otherwise, they catch a fish with only 20% probability. \n",
    "\n",
    "The fisherperson is fishing on the left side. \n",
    "\n",
    "1) Figure out each of the following:\n",
    "- probability of catching a fish given that the school of fish is on the left side, $P(m = catch\\text{ } fish | s = left )$\n",
    "- probability of not catching a fish given that the school of fish is on the left side, $P(m = no \\text{ } fish | s = left)$\n",
    "- probability of catching a fish given that the school of fish is on the right side, $P(m = catch  \\text{ } fish | s = right)$\n",
    "- probability of not catching a fish given that the school of fish is on the right side, $P(m = no \\text{ } fish | s = right)$\n",
    "\n",
    "2) If the fisherperson catches a fish, which side would you guess the school is on? Why?\n",
    "\n",
    "3) If the fisherperson does not catch a fish, which side would you guess the school is on? Why?\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "execution": {
     "iopub.execute_input": "2021-06-02T02:11:14.648051Z",
     "iopub.status.busy": "2021-06-02T02:11:14.647037Z",
     "iopub.status.idle": "2021-06-02T02:11:14.648662Z",
     "shell.execute_reply": "2021-06-02T02:11:14.649059Z"
    }
   },
   "source": [
    "[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D1_BayesianDecisions/solutions/W3D1_Tutorial1_Solution_2284aaf4.py)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the prior exercise, you guessed where the school of fish was based on the measurement you took (watching someone fish). You did this by choosing the state (side of school) that maximized the probability of the measurement. In other words, you estimated the state by maximizing the likelihood (had the highest probability of measurement given state $P(m|s$)). This is called maximum likelihood estimation (MLE) and you've encountered it before during this course, in W1D3!\n",
    "\n",
    "What if you had been going to this river for years and you knew that the fish were almost always on the left side? This would probably affect how you make your estimate - you would rely less on the single new measurement and more on your prior knowledge. This is the idea behind Bayesian inference, as we will see later in this tutorial!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Section 4: Correlation and marginalization\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 517
    },
    "execution": {
     "iopub.execute_input": "2021-06-02T02:11:14.653836Z",
     "iopub.status.busy": "2021-06-02T02:11:14.653358Z",
     "iopub.status.idle": "2021-06-02T02:11:14.718098Z",
     "shell.execute_reply": "2021-06-02T02:11:14.718575Z"
    },
    "outputId": "e4c1eecf-211f-4836-e4ed-aa45c607d0cd"
   },
   "outputs": [],
   "source": [
    "#@title Video 5: Correlation and marginalization\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id='vsDjtWi-BVo', width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we are going to take a step back for a bit and think more generally about the amount of information shared between two random variables. We want to know how much information you gain when you observe one variable (take a measurement) if you know something about another. We will see that the fundamental concept is the same if we think about two attributes, for example the size and color of the fish, or the prior information and the likelihood."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Math Exercise 4: Computing marginal likelihoods\n",
    "\n",
    "To understand the information between two variables, let's first consider the size and color of the fish.\n",
    "\n",
    "| P(X, Y)   | Y = silver   | Y = gold  |\n",
    "| ----------------- |----------|----------|\n",
    "| X = small          | 0.4          | 0.2         |\n",
    "| X = large         | 0.1         |  0.3         |\n",
    "\n",
    "The table above shows us the **joint probabilities**: the probability of both specific attributes occuring together. For example, the probability of a fish being small and silver ($P(X = small, Y = silver$) is 0.4.\n",
    "\n",
    "We want to know what the probability of a fish being small  regardless of color. Since the fish are either silver or gold, this would be the probability of a fish being small and silver plus the probability of a fish being small and gold. This is an example of marginalizing, or averaging out, the variable we are not interested in across the rows or columns.. In math speak: $P(X = small) = \\sum_y{P(X = small, Y)}$. This gives us a **marginal probability**, a probability of a variable outcome (in this case size), regardless of the other variables (in this case color).\n",
    "\n",
    "Please complete the following math problems to further practice thinking through probabilities:\n",
    "\n",
    "1. Calculate the probability of a fish being silver.\n",
    "2. Calculate the probability of a fish being small, large, silver, or gold.\n",
    "3. Calculate the probability of a fish being small OR gold. (Hint: $P(A\\ \\textrm{or}\\ B) = P(A) + P(B) - P(A\\ \\textrm{and}\\ B)$)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "execution": {
     "iopub.execute_input": "2021-06-02T02:11:14.723179Z",
     "iopub.status.busy": "2021-06-02T02:11:14.721950Z",
     "iopub.status.idle": "2021-06-02T02:11:14.723747Z",
     "shell.execute_reply": "2021-06-02T02:11:14.724142Z"
    }
   },
   "source": [
    "[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D1_BayesianDecisions/solutions/W3D1_Tutorial1_Solution_65e69cd1.py)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Think! 4: Covarying probability distributions\n",
    "\n",
    "The relationship between the marginal probabilities and the joint probabilities is determined by the correlation between the two random variables - a normalized measure of how much the variables covary. We can also think of this as gaining some information about one of the variables when we observe a measurement from the other. We will think about this more formally in Tutorial 2. \n",
    "\n",
    "Here, we want to think about how the correlation between size and color of these fish changes how much information we gain about one attribute based on the other. See Bonus Section 1 for the formula for correlation.\n",
    "\n",
    "Use the widget below and answer the following questions:\n",
    "\n",
    "1. When the correlation is zero, $\\rho = 0$, what does the distribution of size tell you about color?\n",
    "2. Set $\\rho$ to something small. As you change the probability of golden fish, what happens to the ratio of size probabilities? Set $\\rho$ larger (can be negative). Can you explain the pattern of changes in the probabilities of size as you change the probability of golden fish?\n",
    "3. Set the probability of golden fish and of large fish to around 65%. As the correlation goes towards 1, how often will you see silver large fish?\n",
    "4. What is increasing the (absolute) correlation telling you about how likely you are to see one of the properties if you see a fish with the other?\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-02T02:11:14.804158Z",
     "iopub.status.busy": "2021-06-02T02:11:14.739035Z",
     "iopub.status.idle": "2021-06-02T02:11:14.978984Z",
     "shell.execute_reply": "2021-06-02T02:11:14.977835Z"
    }
   },
   "outputs": [],
   "source": [
    "# @markdown Execute this cell to enable the widget\n",
    "style = {'description_width': 'initial'}\n",
    "gs = GridspecLayout(2,2)\n",
    "\n",
    "cor_widget = widgets.FloatSlider(0.0, description='ρ', min=-1, max=1, step=0.01)\n",
    "px_widget = widgets.FloatSlider(0.5, description='p(color=golden)', min=0.01, max=0.99, step=0.01, style=style)\n",
    "py_widget = widgets.FloatSlider(0.5, description='p(size=large)', min=0.01, max=0.99, step=0.01, style=style)\n",
    "gs[0,0] = cor_widget\n",
    "gs[0,1] = px_widget\n",
    "gs[1,0] = py_widget\n",
    "\n",
    "\n",
    "@widgets.interact(\n",
    "    px=px_widget,\n",
    "    py=py_widget,\n",
    "    cor=cor_widget,\n",
    ")\n",
    "def make_corr_plot(px, py, cor):\n",
    "    Cmin, Cmax = compute_cor_range(px, py) #allow correlation values\n",
    "    cor_widget.min, cor_widget.max = Cmin+0.01, Cmax-0.01\n",
    "    if cor_widget.value > Cmax:\n",
    "        cor_widget.value = Cmax\n",
    "    if cor_widget.value < Cmin:\n",
    "        cor_widget.value = Cmin\n",
    "    cor = cor_widget.value\n",
    "    P = compute_marginal(px,py,cor)\n",
    "    # print(P)\n",
    "    fig = plot_joint_probs(P)\n",
    "    plt.show(fig)\n",
    "    plt.close(fig)\n",
    "    return None\n",
    "\n",
    "# gs[1,1] = make_corr_plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "execution": {
     "iopub.execute_input": "2021-06-02T02:11:14.982705Z",
     "iopub.status.busy": "2021-06-02T02:11:14.981766Z",
     "iopub.status.idle": "2021-06-02T02:11:14.983396Z",
     "shell.execute_reply": "2021-06-02T02:11:14.983796Z"
    }
   },
   "source": [
    "[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D1_BayesianDecisions/solutions/W3D1_Tutorial1_Solution_9727d8f3.py)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have just seen how two random variables can be more or less independent. The more correlated, the less independent, and the more shared information. We also learned that we can marginalize to determine the marginal likelihood of a hidden state or to find the marginal probability distribution of two random variables. We are going to now complete our journey towards being fully Bayesian!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Section 5: Bayes' Rule and the Posterior"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Marginalization is going to be used to combine our prior knowlege, which we call the **prior**, and our new information from a measurement, the **likelihood**. Only in this case, the information we gain about the hidden state we are interested in, where the fish are, is based on the relationship between the probabilities of the measurement and our prior. \n",
    "\n",
    "We can now calculate the full posterior distribution for the hidden state ($s$) using Bayes' Rule. As we've seen, the posterior is proportional the the prior times the likelihood. This means that the posterior probability of the hidden state ($s$) given a measurement ($m$) is proportional to the likelihood of the measurement given the state times the prior probability of that state (the marginal likelihood):\n",
    "\n",
    "$$ P(s | m) \\propto P(m | s) P(s)  $$\n",
    "\n",
    "We say proportional to instead of equal because we need to normalize to produce a full probability distribution:\n",
    "\n",
    "$$ P(s | m) = \\frac{P(m | s) P(s)}{P(m)}  $$\n",
    "\n",
    "Normalizing by this $P(m)$ means that our posterior is a complete probability distribution that sums or integrates to 1 appropriately. We now can use this new, complete probability distribution for any future inference or decisions we like! In fact, as we will see tomorrow, we can use it as a new prior! Finally, we often call this probability distribution our beliefs over the hidden states, to emphasize that it is our subjective knowlege about the hidden state.\n",
    "\n",
    "For many complicated cases, like those we might be using to model behavioral or brain inferences, the normalization term can be intractable or extremely complex to calculate. We can be careful to choose probability distributions were we can analytically calculate the posterior probability or numerical approximation is reliable. Better yet, we sometimes don't need to bother with this normalization! The normalization term, $P(m)$, is the probability of the measurement. This does not depend on state so is essentially a constant we can often ignore. We can compare the unnormalized posterior distribution values for different states because how they relate to each other is unchanged when divided by the same constant. We will see how to do this to compare evidence for different hypotheses tomorrow. (It's also used to compare the likelihood of models fit using maximum likelihood estimation, as you did in W1D5.)\n",
    "\n",
    "In this relatively simple example, we can compute the marginal probability $P(m)$ easily by using:\n",
    "$$P(m) = \\sum_s P(m | s) P(s)$$\n",
    "We can then normalize so that we deal with the full posterior distribution.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Math Exercise 5: Calculating a posterior probability\n",
    "\n",
    "Our prior is $p(s = left) = 0.3$ and $p(s = right) = 0.7$. In the video, we learned that the chance of catching a fish given they fish on the same side as the school was 50%. Otherwise, it was 10%. We observe a person fishing on the left side. Our likelihood is: \n",
    "\n",
    "\n",
    "| Likelihood: p(m \\| s) | m = catch fish   | m = no fish  |\n",
    "| ----------------- |----------|----------|\n",
    "| s = left          | 0.5          | 0.5         |\n",
    "| s = right         | 0.1        |  0.9       |\n",
    "\n",
    "\n",
    "Calculate the posterior probability (on paper) that:\n",
    "\n",
    "1. The school is on the left if the fisherperson catches a fish: $p(s = left | m = catch fish)$ (hint: normalize by compute $p(m = catch fish)$)\n",
    "2. The school is on the right if the fisherperson does not catch a fish: $p(s = right | m = no fish)$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "execution": {
     "iopub.execute_input": "2021-06-02T02:11:14.987680Z",
     "iopub.status.busy": "2021-06-02T02:11:14.986740Z",
     "iopub.status.idle": "2021-06-02T02:11:14.988318Z",
     "shell.execute_reply": "2021-06-02T02:11:14.988713Z"
    }
   },
   "source": [
    "[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D1_BayesianDecisions/solutions/W3D1_Tutorial1_Solution_46a5b352.py)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Coding Exercise 5: Computing Posteriors\n",
    "\n",
    "Let's implement our above math to be able to compute posteriors for different priors and likelihood.s\n",
    "\n",
    "As before, our prior is $p(s = left) = 0.3$ and $p(s = right) = 0.7$. In the video, we learned that the chance of catching a fish given they fish on the same side as the school was 50%. Otherwise, it was 10%. We observe a person fishing on the left side. Our likelihood is: \n",
    "\n",
    "\n",
    "| Likelihood: p(m \\| s) | m = catch fish   | m = no fish  |\n",
    "| ----------------- |----------|----------|\n",
    "| s = left          | 0.5          | 0.5         |\n",
    "| s = right         | 0.1        |  0.9       |\n",
    "\n",
    "\n",
    "We want our full posterior to take the same 2 by 2 form. Make sure the outputs match your math answers!\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-02T02:11:14.996502Z",
     "iopub.status.busy": "2021-06-02T02:11:14.995971Z",
     "iopub.status.idle": "2021-06-02T02:11:15.080116Z",
     "shell.execute_reply": "2021-06-02T02:11:15.079591Z"
    }
   },
   "outputs": [],
   "source": [
    "def compute_posterior(likelihood, prior):\n",
    "  \"\"\" Use Bayes' Rule to compute posterior from likelihood and prior\n",
    "\n",
    "  Args:\n",
    "    likelihood (ndarray): i x j array with likelihood probabilities where i is\n",
    "                    number of state options, j is number of measurement options\n",
    "    prior (ndarray): i x 1 array with prior probability of each state\n",
    "\n",
    "  Returns:\n",
    "    ndarray: i x j array with posterior probabilities where i is\n",
    "            number of state options, j is number of measurement options\n",
    "\n",
    "  \"\"\"\n",
    "\n",
    "  #################################################\n",
    "  ## TODO for students ##\n",
    "  # Fill out function and remove\n",
    "  raise NotImplementedError(\"Student exercise: implement compute_posterior\")\n",
    "  #################################################\n",
    "\n",
    "  # Compute unnormalized posterior (likelihood times prior)\n",
    "  posterior = ... # first row is s = left, second row is s = right\n",
    "\n",
    "  # Compute p(m)\n",
    "  p_m = np.sum(posterior, axis = 0)\n",
    "\n",
    "  # Normalize posterior (divide elements by p_m)\n",
    "  posterior /= ...\n",
    "\n",
    "  return posterior\n",
    "\n",
    "# Make prior\n",
    "prior = np.array([0.3, 0.7]).reshape((2, 1)) # first row is s = left, second row is s = right\n",
    "\n",
    "# Make likelihood\n",
    "likelihood = np.array([[0.5, 0.5], [0.1, 0.9]]) # first row is s = left, second row is s = right\n",
    "\n",
    "# Compute posterior\n",
    "posterior = compute_posterior(likelihood, prior)\n",
    "\n",
    "# Visualize\n",
    "with plt.xkcd():\n",
    "  plot_prior_likelihood_posterior(prior, likelihood, posterior)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "execution": {
     "iopub.execute_input": "2021-06-02T02:11:15.085360Z",
     "iopub.status.busy": "2021-06-02T02:11:15.084843Z",
     "iopub.status.idle": "2021-06-02T02:11:15.372416Z",
     "shell.execute_reply": "2021-06-02T02:11:15.371908Z"
    }
   },
   "source": [
    "[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D1_BayesianDecisions/solutions/W3D1_Tutorial1_Solution_042b00b7.py)\n",
    "\n",
    "*Example output:*\n",
    "\n",
    "<img alt='Solution hint' align='left' width=669 height=314 src=https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/tutorials/W3D1_BayesianDecisions/static/W3D1_Tutorial1_Solution_042b00b7_0.png>\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interactive Demo 5: What affects the posterior?\n",
    "\n",
    "Now that we can understand the implementation of *Bayes rule*, let's vary the parameters of the prior and likelihood to see how changing the prior and likelihood affect the posterior. \n",
    "\n",
    "In the demo below, you can change the prior by playing with the slider for $p( s = left)$. You can also change the likelihood by changing the probability of catching a fish given that the school is on the left and the probability of catching a fish given that the school is on the right. The fisherperson you are observing is fishing on the left.\n",
    " \n",
    "\n",
    "1.   Keeping the likelihood constant, when does the prior have the strongest influence over the posterior? Meaning, when does the posterior look most like the prior no matter whether a fish was caught or not?\n",
    "2.   Keeping the likelihood constant, when does the prior exert the weakest influence?  Meaning, when does the posterior look least like the prior and depend most on whether a fish was caught or not?\n",
    "3.  Set the prior probability of the state = left to 0.6 and play with the likelihood. When does the likelihood exert the most influence over the posterior?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-02T02:11:15.384872Z",
     "iopub.status.busy": "2021-06-02T02:11:15.380759Z",
     "iopub.status.idle": "2021-06-02T02:11:15.608732Z",
     "shell.execute_reply": "2021-06-02T02:11:15.609114Z"
    }
   },
   "outputs": [],
   "source": [
    "# @markdown Execute this cell to enable the widget\n",
    "style = {'description_width': 'initial'}\n",
    "ps_widget = widgets.FloatSlider(0.3, description='p(s = left)',\n",
    "                                min=0.01, max=0.99, step=0.01)\n",
    "p_a_s1_widget = widgets.FloatSlider(0.5, description='p(fish | s = left)',\n",
    "                                    min=0.01, max=0.99, step=0.01, style=style)\n",
    "p_a_s0_widget = widgets.FloatSlider(0.1, description='p(fish | s = right)',\n",
    "                                    min=0.01, max=0.99, step=0.01, style=style)\n",
    "observed_widget = widgets.Checkbox(value=False, description='Observed fish (m)',\n",
    "                                 disabled=False, indent=False, layout={'width': 'max-content'})\n",
    "\n",
    "@widgets.interact(\n",
    "    ps=ps_widget,\n",
    "    p_a_s1=p_a_s1_widget,\n",
    "    p_a_s0=p_a_s0_widget,\n",
    "    m_right=observed_widget\n",
    ")\n",
    "def make_prior_likelihood_plot(ps,p_a_s1,p_a_s0,m_right):\n",
    "    fig = plot_prior_likelihood(ps,p_a_s1,p_a_s0,m_right)\n",
    "    plt.show(fig)\n",
    "    plt.close(fig)\n",
    "    return None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "execution": {
     "iopub.execute_input": "2021-06-02T02:11:15.614258Z",
     "iopub.status.busy": "2021-06-02T02:11:15.613130Z",
     "iopub.status.idle": "2021-06-02T02:11:15.614816Z",
     "shell.execute_reply": "2021-06-02T02:11:15.615205Z"
    }
   },
   "source": [
    "[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D1_BayesianDecisions/solutions/W3D1_Tutorial1_Solution_a004d68d.py)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Section 6: Making Bayesian fishing decisions\n",
    "\n",
    "We will explore how to consider the expected utility of an action based on our belief (the posterior distribution) about where we think the fish are. Now we have all the components of a Bayesian decision: our prior information, the likelihood given a measurement, the posterior distribution (belief) and our utility (the gains and losses). This allows us to consider the relationship between the true value of the hidden state, $s$, and what we *expect* to get if we take action, $a$, based on our belief!\n",
    "\n",
    "Let's use the following widget to think about the relationship between these probability distributions and utility function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Think! 6: What is more important, the probabilities or the utilities?\n",
    "\n",
    "We are now going to put everything we've learned together to gain some intuitions for how each of the elements that goes into a Bayesian decision comes together. Remember, the common assumption in neuroscience, psychology, economics, ecology, etc. is that we (humans and animals) are tying to maximize our expected utility.\n",
    "\n",
    "1. Can you find a situation where the expected utility is the same for both actions?\n",
    "2. What is more important for determining the expected utility: the prior or a new measurement (the likelihood)?\n",
    "3. Why is this a normative model?\n",
    "4. Can you think of ways in which this model would need to be extended to describe human or animal behavior?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-02T02:11:15.640642Z",
     "iopub.status.busy": "2021-06-02T02:11:15.632788Z",
     "iopub.status.idle": "2021-06-02T02:11:15.978499Z",
     "shell.execute_reply": "2021-06-02T02:11:15.978061Z"
    }
   },
   "outputs": [],
   "source": [
    "# @markdown Execute this cell to enable the widget\n",
    "style = {'description_width': 'initial'}\n",
    "ps_widget = widgets.FloatSlider(0.3, description='p(s)',\n",
    "                                min=0.01, max=0.99, step=0.01)\n",
    "p_a_s1_widget = widgets.FloatSlider(0.5, description='p(fish | s = left)',\n",
    "                                    min=0.01, max=0.99, step=0.01, style=style)\n",
    "p_a_s0_widget = widgets.FloatSlider(0.1, description='p(fish | s = right)',\n",
    "                                    min=0.01, max=0.99, step=0.01, style=style)\n",
    "observed_widget = widgets.Checkbox(value=False, description='Observed fish (m)',\n",
    "                                 disabled=False, indent=False, layout={'width': 'max-content'})\n",
    "\n",
    "@widgets.interact(\n",
    "    ps=ps_widget,\n",
    "    p_a_s1=p_a_s1_widget,\n",
    "    p_a_s0=p_a_s0_widget,\n",
    "    m_right=observed_widget\n",
    ")\n",
    "def make_prior_likelihood_utility_plot(ps, p_a_s1, p_a_s0,m_right):\n",
    "    fig = plot_prior_likelihood_utility(ps, p_a_s1, p_a_s0,m_right)\n",
    "    plt.show(fig)\n",
    "    plt.close(fig)\n",
    "    return None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "execution": {
     "iopub.execute_input": "2021-06-02T02:11:15.982613Z",
     "iopub.status.busy": "2021-06-02T02:11:15.982144Z",
     "iopub.status.idle": "2021-06-02T02:11:15.983704Z",
     "shell.execute_reply": "2021-06-02T02:11:15.983313Z"
    }
   },
   "source": [
    "[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D1_BayesianDecisions/solutions/W3D1_Tutorial1_Solution_3a230382.py)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Summary\n",
    "\n",
    "In this tutorial, you learned about combining prior information with new measurements to update your knowledge using Bayes Rulem, in the context of a fishing problem. \n",
    "\n",
    "Specifically, we covered:\n",
    "\n",
    "* That the likelihood is the probability of the measurement given some hidden state\n",
    "\n",
    "* That how the prior and likelihood interact to create the posterior, the probability of the hidden state given a measurement, depends on how they covary\n",
    "\n",
    "* That utility is the gain from each action and state pair, and the expected utility for an action is the sum of the utility for all state pairs, weighted by the probability of that state happening. You can then choose the action with highest expected utility.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Bonus"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bonus Section 1: Correlation Formula\n",
    "To understand the way we calculate the correlation, we need to review the definition of covariance and correlation.\n",
    "\n",
    "Covariance:\n",
    "\n",
    "$$\n",
    "cov(X,Y) = \\sigma_{XY} = E[(X - \\mu_{x})(Y - \\mu_{y})] = E[X]E[Y] - \\mu_{x}\\mu_{y}\n",
    "$$\n",
    "\n",
    "Correlation:\n",
    "\n",
    "$$\n",
    "\\rho_{XY} = \\frac{cov(Y,Y)}{\\sqrt{V(X)V(Y)}} = \\frac{\\sigma_{XY}}{\\sigma_{X}\\sigma_{Y}}\n",
    "$$"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [
    "jWmfLbhzBpfz"
   ],
   "include_colab_link": true,
   "name": "W3D1_Tutorial1",
   "provenance": [],
   "toc_visible": true
  },
  "kernel": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "kernelspec": {
   "display_name": "Python 3.7.10 64-bit ('nma': conda)",
   "name": "python3710jvsc74a57bd03e19903e646247cead5404f55ff575624523d45cf244c3f93aaf5fa10367032a"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
